/*
 * Copyright (c) 2025.
 * All rights reserved.
 */

#include "driver.h"
#include "path_filter.h"
#include "ioctl_communicator.h"
#include "registry_monitor.h"
#include "string_utils.h"

BOOLEAN			g_IsSetProcessCallback = FALSE;
BOOLEAN			g_IsSetRegisterCallback = FALSE;
SHORT			g_MsgAsyncFlag = -1;

VOID CreateProcessCallback(
	_Inout_ PEPROCESS process,
	_In_ HANDLE processId,
	_Inout_opt_ PPS_CREATE_NOTIFY_INFO createInfo
);

/// <summary>
/// Handle the message from user
/// </summary>
/// <param name="msgFlag"></param>
/// <returns></returns>
VOID HandleMsg(
	SHORT message)
{
	g_MsgAsyncFlag = message;
	return;
}

VOID DriverUnload(
	_In_ PDRIVER_OBJECT driverObj)
{
	UNREFERENCED_PARAMETER(driverObj);
	DBG_PRINT("Process Filter Driver Unload<<<<<<<<\n");

	IoctlUnload(driverObj);

	if (g_IsSetProcessCallback) {
		PsSetCreateProcessNotifyRoutineEx(CreateProcessCallback, TRUE);
	}
	if (g_IsSetRegisterCallback) {
		UninstallRegistryCallback();
	}
	return;
}

NTSTATUS DriverEntry(
	_In_ PDRIVER_OBJECT driverObj,
	_In_ PUNICODE_STRING registryPath)
{
	DBG_PRINT("Process Filter Driver Entry>>>>>>>>\n");

	UNREFERENCED_PARAMETER(registryPath);
	if (driverObj == NULL) {
		DBG_PRINT("driver object is null\n");
		return STATUS_INVALID_PARAMETER;
	}
	driverObj->DriverUnload = DriverUnload;

	NTSTATUS status;

	status = IoctlLoad(driverObj, HandleMsg);
	if (STATUS_SUCCESS != status) {
		DBG_PRINT("IoctlLoad: Fail\n");
		return status;
	}

	// set create process notify
	status = PsSetCreateProcessNotifyRoutineEx(CreateProcessCallback, FALSE);
	if (STATUS_SUCCESS != status) {
		g_IsSetProcessCallback = FALSE;
		DBG_PRINT("SetCreateProcessNotify: Fail\n");
		return status;
	}

	g_IsSetProcessCallback = TRUE;
	DBG_PRINT("SetCreateProcessNotify: Succeed\n");

	// set register change notify
	status = InstallRegistryCallback(driverObj);
	if (STATUS_SUCCESS != status) {
		g_IsSetRegisterCallback = FALSE;
		DBG_PRINT("InstallRegistryCallback: Fail\n");
		return status;
	}
	g_IsSetRegisterCallback = TRUE;
	DBG_PRINT("InstallRegistryCallback: Succeed\n");

	return STATUS_SUCCESS;
}

VOID CreateProcessCallback(
	_Inout_ PEPROCESS process,
	_In_ HANDLE processId,
	_Inout_opt_ PPS_CREATE_NOTIFY_INFO createInfo)
{
	UNREFERENCED_PARAMETER(process);
	// Process destroy
	if (createInfo == NULL) {
		return;
	}
	PCUNICODE_STRING filePath = createInfo->ImageFileName;
	if (IsNullOrEmpty(filePath)) {
		return;
	}

	PMSG_ITEM	msgItem = NULL;
	BOOLEAN		isAllow = IsPathAllowed(filePath);
	MSG_PARAM   msgParam = { MSG_TYPE_EXE_ASK, filePath, processId, createInfo->ParentProcessId };
	if (g_MsgAsyncFlag == MSG_TYPE_DEBUG_ALL) {
		msgParam.type = MSG_TYPE_DEBUG_ALL;
		if (IoctlSendMsg(msgParam, TRUE, &msgItem) && msgItem != NULL) {
			SetMsgState(msgItem, MSG_STATE_IDLE);
		}
		return;
	}
	if (isAllow) {
		return;
	}
	if (IoctlSendMsg(msgParam, TRUE, &msgItem) && msgItem != NULL) {
		if (msgItem->Message == MSG_TYPE_EXE_DENY) {
			createInfo->CreationStatus = STATUS_ACCESS_DISABLED_NO_SAFER_UI_BY_POLICY;
		}
		SetMsgState(msgItem, MSG_STATE_IDLE);
	}
	return;
}